# -*- coding: utf-8 -*-

'''
Copyright: 
Descripttion: 
version: 
Author: chengx
Date: 2021-03-11 14:02:45
LastEditors: chengx
LastEditTime: 2021-05-28 22:31:48
'''

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import copy
import bisect
import random

from tqdm import tqdm


class ABC:
    """
    Parameters: data        : 高光谱数据
                bandAmount  : 选择的目标波段数
                scope       : 每个类里波段的限制范围
    Description:
    Returns:
    """
    def __init__(self,data,bandAmount,scope):

        #初始化各参数
        self.n_bands = bandAmount #选择的波段数
        self.scope = scope # 选择边界
        #蜂群算法基本参数
        self.iter_max = 60 #最大循环次数
        self.e_size = 100  #雇佣峰数量
        self.l_size = 100 #观察蜂数目
        self.L = np.around(0.6*self.n_bands*self.e_size) #蜜源抛弃上界

        #创建各记录矩阵
        self.pops = np.zeros([self.e_size,self.n_bands]) #蜜源位置矩阵
        self.popCost = np.zeros(self.e_size) #蜜源代价矩阵popCost(e_size)
        self.Probability = np.zeros([self.e_size,1]) #蜜源选择概率矩阵Probability(e_size)
        self.bestSol = np.zeros([self.iter_max+1,self.n_bands]) #历史最优蜜源矩阵bestSol(iter_max+1行n_bands列，保存每次迭代中的最优蜜源位置)
        self.mineCount = np.zeros([self.e_size,1]) #蜜源开采次数矩阵mineCount(e_size)
        self.bestCost = np.ones(self.iter_max+1)
        self.data = data

    def initParams(self):
        #初始化蜜源位置
        temp = copy.deepcopy(self.scope)
        for i in range(self.e_size):
            initPos = [random.sample(list(temp[k]), 1)[0] for k in range(len(temp))]
            self.pops[i,:] = np.array(initPos)
        del temp
        # 初始化种群
        for i in range(self.e_size):
            self.popCost[i] = CostFunction(self.pops[i],self.data)
            if self.popCost[i] > self.bestCost[0]:
                self.bestCost[0] = self.popCost[i]
                self.bestSol[0] = self.pops[i]

    def run(self):
        # for iter in range(self.iter_max):
        for iter in tqdm(range(1,self.iter_max)):
            #雇佣蜂阶段
            for i in range(self.e_size):
                while True:
                    k = np.random.randint(0,self.e_size)
                    if k != i:
                        break
                phi = np.random.rand()
                NewPosition = self.pops[i] + np.ceil(phi*(self.pops[i]-self.pops[k]))

                # 边界处理
                NewPosition = self.clusterScoped( NewPosition )

                #进行贪婪选择
                NewCost = CostFunction(NewPosition,self.data) #计算新蜜源的代价值，在当前的和新的中进行贪婪选择
                if NewCost > self.popCost[i]:
                    self.pops[i] = NewPosition
                    self.popCost[i] = NewCost
                else:
                    self.mineCount[i][0] = self.mineCount[i][0]+1

            #计算选择概率矩阵
            Mean = np.mean(self.popCost)
            for i in range(self.e_size):
                self.Probability[i][0] = np.exp(-self.popCost[i]/Mean)

            self.Probability = self.Probability/np.sum(self.Probability)
            P = np.cumsum(self.Probability)

            # 观察蜂阶段
            for k in range(self.l_size):
                r = np.random.rand()
                j =  bisect.bisect(P, r)
                #重复雇佣蜂操作
                while True:
                    k = np.random.randint(0,self.e_size)
                    if k != j:
                        break
                phi = np.random.rand() 
                NewPosition = self.pops[i] + np.ceil(phi*(self.pops[i]-self.pops[k]))

                # 边界处理
                NewPosition = self.clusterScoped(NewPosition)

                #进行贪婪选择
                NewCost = CostFunction(NewPosition,self.data)
                if NewCost > self.popCost[i]:
                    self.pops[i] = NewPosition
                    self.popCost[i] = NewCost
                else:
                    self.mineCount[i][0] = self.mineCount[i][0]+1
                # print('{}:{}   bands index{}'.format('looker-bee',k,self.pops[i])) 
            #侦查蜂阶段
            for i in range(self.e_size):
                if self.mineCount[i][0] >= self.L:
                    self.pops[i] = np.ceil(8*np.random.rand(1,self.n_bands)) - 4
                    self.popCost[i] = CostFunction(self.pops[i],self.data)
                    self.mineCount[i][0] = 0

            #保存历史最优解
            #遍历所有蜜源节点，将最优节点保存在历史最优蜜源矩阵bestSol的第iter+1行中(第一行是第0代，初始值)。
            for i in range(self.e_size):
                if self.popCost[i] > self.bestCost[iter+1]:
                    self.bestCost[iter+1] = self.popCost[i]
                    self.bestSol[iter+1] = self.pops[i]

        return self.bestSol[iter+1]
        
    def clusterScoped(self,NewPosition):
        # 边界处理
        for i in range(self.n_bands):
            if NewPosition[i] in self.scope[i]:
                pass
            else:#边界处理可以有不同的方式
                # NewPosition[i] = np.min(self.scope[i])
                NewPosition[i] = random.sample(list(self.scope[i]), 1)[0]

        return NewPosition


def CostFunction(input,dataHSI,index='OIF'):
    """
    Parameters: input   波段号
                dataHSI 光谱数据 (格式: numpy,二维,point*bands)
                index   选择优化函数索引
    Description: 这个就是目标函数，优化的方向
    Returns: 函数值
    """
    input = input.astype(int)
    X=dataHSI[input,:] # 调出选择的波段数据

    s = 0
    r = 0
    if index == 'correlation':
        return 0
    if index == 'OIF':# 越大越好，标准差要大，相关系数要小
        for i in range(input.shape[0]):
            s +=np.std(X[:,input[i]-1])# 标准差
        
        X = X.T
        X = pd.DataFrame(X)
        r = X.corr()
        
        rSum = r.values.sum()
        rSum = (rSum - 12)/2# 其实这里可以不做，对结果无影响
        oif = s/rSum

        # print('oif',oif)
        return oif
    